package com.xc.find;

import lombok.Data;

@Data
public class TreeNode {
    int key;
    boolean color;
    TreeNode left;
    TreeNode right;

    public TreeNode(int key, boolean color) {
        // 这里需要一个构造函数，用于创建一个带有颜色的节点
        this.key = key;
        this.color = color;
    }

    TreeNode(int key) {
        this.key = key;
        this.color = true;
        this.left = null;
        this.right = null;
    }
}


class RedBlackTree {
    private TreeNode root;

    public RedBlackTree() {
        root = null;
    }

    public void insert(int key) {
        if (root == null) {
            root = new TreeNode(key, true);
        } else {
            TreeNode current = root;
            while (true) {
                if (key < current.key) {
                    if (current.left == null) {
                        current.left = new TreeNode(key, true);
                        fixInsert(root);
                        break;
                    }
                    current = current.left;
                } else {
                    if (current.right == null) {
                        current.right = new TreeNode(key, false);
                        fixInsert(root);
                        break;
                    }
                    current = current.right;
                }
            }
        }
    }

    private void fixInsert(TreeNode node) {
        if (node == null) return;
        if (isRed(node.right) && !isRed(node.left)) {
            rotateLeft(node);
            fixInsert(node.left);
        } else if (isRed(node.left) && !isRed(node.right)) {
            rotateRight(node);
            fixInsert(node.right);
        } else if (isRed(node.left) && isRed(node.right)) {
            flipColors(node);
        }
    }

    private boolean isRed(TreeNode node) {
        if (node == null) return false;
        return node.color;
    }

    private void rotateLeft(TreeNode node) {
        TreeNode temp = node.right;
        node.right = temp.left;
        temp.left = node;
        temp.color = node.color;
        node.color = true;
    }

    private void rotateRight(TreeNode node) {
        TreeNode temp = node.left;
        node.left = temp.right;
        temp.right = node;
        temp.color = node.color;
        node.color = true;
    }

    private void flipColors(TreeNode node) {
        node.color = true;
        node.left.color = false;
        node.right.color = false;
    }

    public boolean contains(int key) {
        TreeNode current = root;
        while (current != null) {
            if (key < current.key) {
                current = current.left;
            } else if (key > current.key) {
                current = current.right;
            } else {
                return true;
            }
        }
        return false;
    }
}